"""
大模型生成模块

这个模块负责：
1. 连接ollama服务
2. 调用deepseek-r1:14b模型
3. 构造提示词
4. 生成回答

作者：RAG学习项目
"""

import requests
import json
import logging
from typing import Dict, Any, Optional, List
from .config import config

# 配置日志
logging.basicConfig(level=getattr(logging, config.LOG_LEVEL))
logger = logging.getLogger(__name__)

class LLMClient:
    """
    大模型客户端类
    
    负责与ollama服务器通信，调用深度求索模型
    """
    
    def __init__(self, host: str = None, port: int = None, model: str = None):
        """
        初始化LLM客户端
        
        Args:
            host (str, optional): ollama服务器地址，默认使用配置值
            port (int, optional): ollama服务器端口，默认使用配置值  
            model (str, optional): 使用的模型名称，默认使用配置值
        """
        self.host = host or config.OLLAMA_HOST
        self.port = port or config.OLLAMA_PORT
        self.model = model or config.OLLAMA_MODEL
        self.base_url = f"http://{self.host}:{self.port}"
        
        logger.info(f"初始化LLM客户端，服务器: {self.base_url}, 模型: {self.model}")
    
    def test_connection(self) -> Dict[str, Any]:
        """
        测试与ollama服务器的连接
        
        Returns:
            Dict[str, Any]: 连接测试结果
        """
        logger.info("测试ollama服务器连接...")
        
        try:
            # 测试服务器是否可达
            response = requests.get(f"{self.base_url}/api/tags", timeout=10)
            
            if response.status_code == 200:
                data = response.json()
                models = [model['name'] for model in data.get('models', [])]
                
                result = {
                    'success': True,
                    'message': '服务器连接成功',
                    'available_models': models,
                    'target_model_available': self.model in models
                }
                
                if self.model in models:
                    logger.info(f"✅ 连接成功，目标模型 {self.model} 可用")
                else:
                    logger.warning(f"⚠️ 连接成功，但目标模型 {self.model} 不可用")
                    logger.info(f"可用模型: {models}")
                
                return result
            else:
                result = {
                    'success': False,
                    'message': f'服务器响应错误: {response.status_code}',
                    'available_models': [],
                    'target_model_available': False
                }
                logger.error(f"❌ 服务器响应错误: {response.status_code}")
                return result
                
        except requests.exceptions.ConnectTimeout:
            result = {
                'success': False,
                'message': '连接超时，请检查网络和服务器地址',
                'available_models': [],
                'target_model_available': False
            }
            logger.error("❌ 连接超时")
            return result
            
        except requests.exceptions.ConnectionError:
            result = {
                'success': False,
                'message': '连接失败，请检查服务器是否运行',
                'available_models': [],
                'target_model_available': False
            }
            logger.error("❌ 连接失败")
            return result
            
        except Exception as e:
            result = {
                'success': False,
                'message': f'未知错误: {str(e)}',
                'available_models': [],
                'target_model_available': False
            }
            logger.error(f"❌ 未知错误: {str(e)}")
            return result
    
    def generate_response(self, prompt: str, max_tokens: int = 1000, temperature: float = 0.7) -> Dict[str, Any]:
        """
        使用LLM生成回答
        
        Args:
            prompt (str): 输入提示词
            max_tokens (int): 最大生成token数
            temperature (float): 温度参数，控制生成的随机性
            
        Returns:
            Dict[str, Any]: 生成结果
        """
        logger.info(f"开始生成回答，提示词长度: {len(prompt)} 字符")
        
        try:
            # 构造请求数据
            request_data = {
                "model": self.model,
                "prompt": prompt,
                "stream": False,  # 不使用流式响应，方便处理
                "options": {
                    "num_predict": max_tokens,
                    "temperature": temperature,
                    "top_p": 0.9,
                    "top_k": 40
                }
            }
            
            # 发送请求
            response = requests.post(
                f"{self.base_url}/api/generate",
                json=request_data,
                timeout=60  # 生成可能需要较长时间
            )
            
            if response.status_code == 200:
                data = response.json()
                generated_text = data.get('response', '').strip()
                
                result = {
                    'success': True,
                    'response': generated_text,
                    'model_used': data.get('model', self.model),
                    'total_duration': data.get('total_duration', 0),
                    'load_duration': data.get('load_duration', 0),
                    'prompt_eval_count': data.get('prompt_eval_count', 0),
                    'eval_count': data.get('eval_count', 0)
                }
                
                logger.info(f"✅ 生成成功，回答长度: {len(generated_text)} 字符")
                return result
            else:
                result = {
                    'success': False,
                    'response': '',
                    'error': f'请求失败: {response.status_code}',
                    'details': response.text
                }
                logger.error(f"❌ 生成失败: {response.status_code}")
                return result
                
        except requests.exceptions.Timeout:
            result = {
                'success': False,
                'response': '',
                'error': '请求超时，模型生成时间过长'
            }
            logger.error("❌ 请求超时")
            return result
            
        except Exception as e:
            result = {
                'success': False,
                'response': '',
                'error': f'生成失败: {str(e)}'
            }
            logger.error(f"❌ 生成失败: {str(e)}")
            return result

class PromptTemplate:
    """
    提示词模板类
    
    负责构造各种场景下的提示词
    """
    
    def __init__(self):
        """初始化提示词模板"""
        logger.info("初始化提示词模板")
    
    def create_rag_prompt(self, question: str, context: str, max_context_length: int = 2000) -> str:
        """
        创建RAG问答的提示词
        
        Args:
            question (str): 用户问题
            context (str): 检索到的相关上下文
            max_context_length (int): 上下文最大长度
            
        Returns:
            str: 完整的提示词
        """
        # 限制上下文长度，避免超出模型限制
        if len(context) > max_context_length:
            context = context[:max_context_length] + "..."
            logger.info(f"上下文过长，已截断到 {max_context_length} 字符")
        
        prompt = f"""你是一个专业的家谱知识问答助手，请根据提供的翁氏宗谱相关资料来回答用户的问题。

**重要指示：**
1. 请仔细阅读提供的参考资料
2. 只基于参考资料中的信息来回答问题
3. 特别注意：资料中的"金钟人"和"族姓赞"是诗文的标题名称
4. 当用户询问诗文、作品、标题时，请重点关注这些关键词
5. 如果参考资料中没有相关信息，请明确说明"根据提供的资料无法找到相关信息"
6. 回答要准确、详细且易于理解
7. 如果涉及人名、地名、年代等具体信息，请保持准确

**参考资料：**
{context}

**用户问题：**
{question}

**特别提醒：如果用户询问翁和平的诗文或作品标题，请仔细查看参考资料中是否包含"金钟人"、"族姓赞"等关键词，这些很可能就是诗文的标题。

**回答：**"""
        
        logger.debug(f"创建RAG提示词，问题: {question[:50]}..., 上下文长度: {len(context)}")
        return prompt
    
    def create_simple_prompt(self, question: str) -> str:
        """
        创建简单对话的提示词（不使用RAG）
        
        Args:
            question (str): 用户问题
            
        Returns:
            str: 简单的提示词
        """
        prompt = f"""你是一个友好的AI助手。请回答用户的问题。

用户问题：{question}

回答："""
        
        logger.debug(f"创建简单提示词，问题: {question[:50]}...")
        return prompt

class ResponseGenerator:
    """
    回答生成器类
    
    整合LLM客户端和提示词模板，提供完整的回答生成功能
    """
    
    def __init__(self):
        """初始化回答生成器"""
        self.llm_client = LLMClient()
        self.prompt_template = PromptTemplate()
        
        logger.info("初始化回答生成器")
    
    def generate_rag_answer(self, question: str, context: str) -> Dict[str, Any]:
        """
        基于RAG生成回答
        
        Args:
            question (str): 用户问题
            context (str): 检索到的上下文
            
        Returns:
            Dict[str, Any]: 生成结果
        """
        logger.info(f"生成RAG回答，问题: {question[:50]}...")
        
        try:
            # 创建提示词
            prompt = self.prompt_template.create_rag_prompt(question, context)
            
            # 生成回答
            result = self.llm_client.generate_response(prompt)
            
            if result['success']:
                logger.info("✅ RAG回答生成成功")
            else:
                logger.error(f"❌ RAG回答生成失败: {result.get('error', '未知错误')}")
            
            return result
            
        except Exception as e:
            logger.error(f"❌ RAG回答生成异常: {str(e)}")
            return {
                'success': False,
                'response': '',
                'error': f'生成异常: {str(e)}'
            }
    
    def generate_simple_answer(self, question: str) -> Dict[str, Any]:
        """
        生成简单回答（不使用RAG）
        
        Args:
            question (str): 用户问题
            
        Returns:
            Dict[str, Any]: 生成结果
        """
        logger.info(f"生成简单回答，问题: {question[:50]}...")
        
        try:
            # 创建简单提示词
            prompt = self.prompt_template.create_simple_prompt(question)
            
            # 生成回答
            result = self.llm_client.generate_response(prompt)
            
            if result['success']:
                logger.info("✅ 简单回答生成成功")
            else:
                logger.error(f"❌ 简单回答生成失败: {result.get('error', '未知错误')}")
            
            return result
            
        except Exception as e:
            logger.error(f"❌ 简单回答生成异常: {str(e)}")
            return {
                'success': False,
                'response': '',
                'error': f'生成异常: {str(e)}'
            }

# 测试函数
def test_llm_connection():
    """
    测试LLM连接功能
    """
    print("=== 测试LLM连接 ===")
    
    try:
        # 创建LLM客户端
        client = LLMClient()
        
        # 测试连接
        connection_result = client.test_connection()
        print(f"连接测试结果: {connection_result}")
        
        if connection_result['success']:
            # 测试简单生成
            test_prompt = "你好，请用中文简单介绍一下自己。"
            generation_result = client.generate_response(test_prompt)
            
            if generation_result['success']:
                print(f"生成测试成功:")
                print(f"回答: {generation_result['response']}")
                
                # 测试回答生成器
                generator = ResponseGenerator()
                simple_result = generator.generate_simple_answer("你能帮助我了解家谱信息吗？")
                
                if simple_result['success']:
                    print(f"回答生成器测试成功:")
                    print(f"回答: {simple_result['response']}")
                else:
                    print(f"回答生成器测试失败: {simple_result.get('error')}")
            else:
                print(f"生成测试失败: {generation_result.get('error')}")
        
        print("✅ LLM连接测试完成")
        
    except Exception as e:
        print(f"❌ LLM连接测试失败: {e}")

if __name__ == "__main__":
    test_llm_connection() 